import numpy as np
import matplotlib.pyplot as plt

# 定义 SH 曲线函数
def SH1_CURVE(t, x4):
    if t <= 0.0:
        return 0.0
    elif t < x4:
        return (t / x4) ** 2.5
    else:
        return 1.0

def SH2_CURVE(t, x4):
    if t <= 0.0:
        return 0.0
    elif t <= x4:
        return 0.5 * (t / x4) ** 2.5
    elif t < 2 * x4:
        return 1 - 0.5 * (2 - t / x4) ** 2.5
    else:
        return 1.0

# 读取参数和数据
# para = np.loadtxt('GR4J_Parameter.txt')  # 该文件包含占位符，使用默认值
other_para = np.loadtxt('others.txt')
data = np.loadtxt('inputData.txt')

area = other_para[0]
upperTankRatio = other_para[1]
lowerTankRatio = other_para[2]

P = data[:, 0]
E = data[:, 1]
Qobs = data[:, 2]
Qobs_mm = Qobs * 86.4 / area  # m³/s -> mm/d
nStep = data.shape[0]

# 用于记录最优结果
max_nse = -np.inf
best_params = None
best_Q = None

# 遍历参数组合 - 根据GR4J标准参数范围优化
print(f"开始参数优化，数据长度: {nStep}, 水文流域面积: {area}")

count = 0

# Set parameter search space according to GR4J standard ranges
x1_range = range(100, 1201, 200)    # x1: 100-1200, step 200
x2_range = range(-50, 31, 10)       # x2: -5 to 3, step 1.0  
x3_range = range(20, 301, 20)       # x3: 20-300, step 40
x4_range = range(1, 71, 10)         # x4: 0.1-7.0, step 1.0

total_combinations = len(x1_range) * len(x2_range) * len(x3_range) * len(x4_range)
print(f"Total parameter combinations to test: {total_combinations}")

for x1_raw in x1_range:
    x1 = float(x1_raw)
    for x2_raw in x2_range:
        for x3 in x3_range:
            for x4_raw in x4_range:
                x2 = x2_raw / 10.0
                x4 = x4_raw / 10.0
                count += 1
                if count % 100 == 0:
                    print(f"已完成 {count}/{total_combinations} 组合 (当前最佳NSE: {max_nse:.4f})")


                # 初始化
                Pn = np.zeros(nStep)
                En = np.zeros(nStep)
                Ps = np.zeros(nStep)
                Es = np.zeros(nStep)
                Perc = np.zeros(nStep)
                Pr = np.zeros(nStep)
                Qr = np.zeros(nStep)
                Qd = np.zeros(nStep)
                Q = np.zeros(nStep)

                maxDayDelay = 10
                SH1 = np.array([SH1_CURVE(i + 1, x4) for i in range(maxDayDelay)])
                SH2 = np.array([SH2_CURVE(i + 1, x4) for i in range(2 * maxDayDelay)])
                UH1 = np.diff(np.concatenate(([0], SH1)))
                UH2 = np.diff(np.concatenate(([0], SH2)))

                for i in range(nStep):
                    if P[i] >= E[i]:
                        Pn[i] = P[i] - E[i]
                        En[i] = 0
                    else:
                        Pn[i] = 0
                        En[i] = E[i] - P[i]

                S0 = upperTankRatio * x1
                R0 = lowerTankRatio * x3
                S = np.zeros(nStep)
                R = np.zeros(nStep)
                UH_Fast = np.zeros((nStep, maxDayDelay))
                UH_Slow = np.zeros((nStep, 2 * maxDayDelay))
                S_TEMP = S0
                R_TEMP = R0

                for i in range(nStep):
                    S[i] = S_TEMP
                    R[i] = R_TEMP

                    if Pn[i] != 0:
                        Ps[i] = x1 * (1 - (S[i] / x1) ** 2) * np.tanh(Pn[i] / x1) / (1 + S[i] / x1 * np.tanh(Pn[i] / x1))
                        Es[i] = 0
                    elif En[i] != 0:
                        Ps[i] = 0
                        Es[i] = (S[i] * (2 - (S[i] / x1)) * np.tanh(En[i] / x1)) / (1 + (1 - S[i] / x1) * np.tanh(En[i] / x1))

                    S_TEMP = S[i] - Es[i] + Ps[i]
                    Perc[i] = S_TEMP * (1 - (1 + (4.0 / 9.0 * (S_TEMP / x1)) ** 4) ** (-0.25))
                    Pr[i] = Perc[i] + (Pn[i] - Ps[i])
                    S_TEMP -= Perc[i]

                    F = x2 * (R[i] / x3) ** 3.5
                    R_Fast = Pr[i] * 0.9
                    R_Slow = Pr[i] * 0.1

                    UH_Fast[i, :] = R_Fast * UH1
                    UH_Slow[i, :] = R_Slow * UH2

                    if i > 0:
                        UH_Fast[i, :-1] += UH_Fast[i - 1, 1:]
                        UH_Slow[i, :-1] += UH_Slow[i - 1, 1:]

                    R_TEMP = max(0, R_TEMP + UH_Fast[i, 0] + F)
                    Qr[i] = R_TEMP * (1 - (1 + (R_TEMP / x3) ** 4) ** (-0.25))
                    R_TEMP -= Qr[i]
                    Qd[i] = max(0, UH_Slow[i, 0] + F)
                    Q[i] = Qr[i] + Qd[i]

                # 模型评估（跳过前365天）
                Qobs_eval = Qobs_mm[365:]
                Q_eval = Q[365:]
                if len(Qobs_eval) > 0:
                    Q_mean = np.mean(Qobs_eval)
                    numerator = np.sum((Qobs_eval - Q_eval) ** 2)
                    denominator = np.sum((Qobs_eval - Q_mean) ** 2)
                    NSE = 1 - numerator / denominator if denominator != 0 else -np.inf

                    if NSE > max_nse:
                        max_nse = NSE
                        best_params = (x1, x2, x3, x4)
                        best_Q = Q.copy()

print(f"最优参数组合: x1={best_params[0]}, x2={best_params[1]}, x3={best_params[2]}, x4={best_params[3]}")
print(f"最优NSE值: {max_nse:.4f}")

# Plot results
plt.figure(figsize=(10, 5))
plt.plot(range(nStep), Qobs_mm, label="Observed", color="black")
plt.plot(range(nStep), best_Q, label="Simulated", color="red", linestyle='--')
plt.title(f"GR4J Model Simulation Results (NSE={max_nse:.4f})")
plt.xlabel("Time (days)")
plt.ylabel("Flow (mm/d)")
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
